# -*- coding: utf-8 -*-

'''对已经分词完毕的.txt文件进行生成词云.png图片'''

from PIL import Image
from os import path
import numpy as np
import matplotlib.pyplot as plt
from wordcloud import WordCloud, STOPWORDS, ImageColorGenerator


def create_wordcloud():

    '''生成词云图片'''

    d = path.dirname(__file__)

    # Read the whole text.
    text = open(path.join(d + '/BaiduSpider_Result/' + '0001.txt'), 'r', encoding='utf-8').read()

    # read the mask image
    alice_mask = np.array(Image.open(path.join(d, "view3.jpg")))

    # The forbidden words
    stopwords = set(STOPWORDS)
    stopwords.add("said")

    wc = WordCloud(
        font_path="simkai.ttf",
        background_color="white",
        max_words=5000,
        mask=alice_mask,  # 设置背景图片
        stopwords=stopwords  # 字体排列x,y轴运动
        )

    # generate word cloud 生成词云, 可以用generate输入全部文本(中文不好分词)
    wc.generate(text)

    # store to file
    '''可以从随机图形库中选择生成的词云图片'''
    wc.to_file(path.join(d, "view4.png"))

    # show
    plt.imshow(wc, interpolation='bilinear')
    plt.axis("off")
    plt.show()

if __name__ == '__main__':
    create_wordcloud()
